
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <thrust/async/for_each.h>
#include <thrust/async/transform.h>
#include <thrust/device_ptr.h>
#include <thrust/device_vector.h>
#include <thrust/extrema.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
#include <thrust/type_traits/integer_sequence.h>
#include <vector_functions.h>

#include <algorithm/analytical_model/algorithm/cpu/ShadowBlockCalculator.h>
#include <algorithm/analytical_model/algorithm/cuda/GeometryUtils.cuh>
#include <algorithm/analytical_model/algorithm/cuda/Math.cuh>
#include <algorithm/analytical_model/algorithm/cuda/ShadowBlockCalculator.cuh>
#include <algorithm/analytical_model/algorithm/cuda/ShadowBlockCollision.cuh>
#include <algorithm/analytical_model/algorithm/cuda/Utils.cuh>
#include <algorithm/analytical_model/algorithm/cuda/VectorMath.cuh>
#include <algorithm/analytical_model/algorithm/cuda/grid/RectangleGrid.cuh>
#include <utils/BitUtils.h>
#include <utils/Format.h>
#include <utils/Print.h>
#include <utils/Utils.h>
#include <utils/config.h>

#include <cstddef>
#include <memory>
#include <tuple>

using namespace solar;
using namespace solar::cuda;

ShadowBlockCalculator::ShadowBlockCalculator() { cudaStreamCreate(&stream_); }

ShadowBlockCalculator::~ShadowBlockCalculator()
{
    if (d_verts_ != nullptr)
    {
        CHECK_CUDA_ERRORS(cudaFreeAsync(d_verts_, stream_));
    }
    if (d_sb_result_ != nullptr)
    {
        CHECK_CUDA_ERRORS(cudaFreeAsync(d_sb_result_, stream_));
    }
    if (d_microhelio_ != nullptr)
    {
        CHECK_CUDA_ERRORS(cudaFreeAsync(d_microhelio_, stream_));
    }
    cudaStreamDestroy(stream_);
}

void ShadowBlockCalculator::init(std::vector<inter_op::Float3>& input_vertices)
{
    n_helio_ = input_vertices.size() / 4;
    CHECK_CUDA_ERRORS(cudaHostRegister(input_vertices.data(), sizeof(float3) * n_helio_ * 4,
                                       cudaHostRegisterMapped));
    float3* d_input_vertices_ptr = nullptr;
    CHECK_CUDA_ERRORS(
        cudaHostGetDevicePointer(std::addressof(d_input_vertices_ptr), input_vertices.data(), 0));
    CHECK_CUDA_ERRORS(
        cudaMallocAsync(std::addressof(this->d_verts_), sizeof(float3) * 3 * n_helio_, stream_));
    auto* d_verts_ptr = this->d_verts_;
    thrust::for_each(thrust::device.on(stream_), thrust::make_counting_iterator(0),
                     thrust::make_counting_iterator(n_helio_),
                     [d_input_vertices_ptr, d_verts_ptr] __host__ __device__(const int& index)
                     {
                         std::memcpy(d_verts_ptr + static_cast<ptrdiff_t>(3 * index),
                                     d_input_vertices_ptr + static_cast<ptrdiff_t>(4 * index),
                                     sizeof(float3) * 3);
                     });
    CHECK_CUDA_ERRORS(cudaHostUnregister(input_vertices.data()));
}

auto ShadowBlockCalculator::getResult() const -> float4* { return d_sb_result_; }

auto ShadowBlockCalculator::getMicroHeliostatShadowBlock() const -> bool* { return d_microhelio_; }

void ShadowBlockCalculator::setHelioRowCol(const int2& row_col) { helio_row_col_ = row_col; }

void ShadowBlockCalculator::copyHeliostatVerticesToGPU()
{
    // auto* h_verts = vertices_.data();
    // cpu2gpuAsync(d_verts_, h_verts, 3 * n_helio_, stream_);
}

void ShadowBlockCalculator::initResultArray()
{
    CHECK_CUDA_ERRORS(
        cudaMallocAsync(std::addressof(d_sb_result_), n_helio_ * sizeof(float4), stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(d_sb_result_, 0, n_helio_ * sizeof(float4), stream_));

    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(d_microhelio_),
                                      sizeof(bool) * n_helio_ * helio_row_col_.x * helio_row_col_.y,
                                      stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(
        d_microhelio_, 0, sizeof(bool) * n_helio_ * helio_row_col_.x * helio_row_col_.y, stream_));

    CHECK_CUDA_ERRORS(cudaMallocAsync(std::addressof(d_microhelio_per_row_),
                                      sizeof(bool) * n_helio_ * helio_row_col_.x, stream_));
    CHECK_CUDA_ERRORS(cudaMemsetAsync(d_microhelio_per_row_, 0,
                                      sizeof(bool) * n_helio_ * helio_row_col_.x, stream_));
}

__global__ static void shadowBlocking(RectangleGrid rectangle_grid, float3* d_verts,
                                      float4* d_sb_result, bool* d_microhelio,
                                      bool* d_microhelio_per_row, int2 row_col, int n,
                                      float3 sun_dir)
{
    long long my_id = getThreadID();
    if (my_id >= row_col.x * row_col.y * n)
    {
        return;
    }
    const int row_col_mul = row_col.x * row_col.y;
    int helio_id = my_id / (row_col_mul);
    int micro_helio_id = my_id % (row_col_mul);
    //    printf("%d\n",micro_helio_id);
    const auto& vert0 = d_verts[helio_id * 3];
    const auto& vert1 = d_verts[helio_id * 3 + 1];
    const auto& vert2 = d_verts[helio_id * 3 + 2];
    const auto vec1 = vert1 - vert0;
    const auto vec2 = vert2 - vert1;
    float3 normal = normalize(cross(vec2, vec1));
    // if (dot(normal, sun_dir) > 0)
    // {
    //     normal = -normal;
    // }
    normal -= 2 * normal * float(dot(normal, sun_dir) > 0);
    float3 origin = (0.5F + fdividef(micro_helio_id, row_col.y)) * (vec2) / row_col.x +
                    (0.5F + micro_helio_id % row_col.y) * (vec1) / row_col.y + vert0;
    //    printf("%f %f %f\n",origin.x,origin.y,origin.z);
    float3 reflect_dir = normalize(reflect(sun_dir, normal));
    sun_dir = -sun_dir;
    if (collision(origin, sun_dir, &rectangle_grid, d_verts, helio_id) ||
        collision(origin, reflect_dir, &rectangle_grid, d_verts, helio_id))
    {
        d_microhelio[micro_helio_id + helio_id * row_col.x * row_col.y] = true;
        d_microhelio_per_row[helio_id * row_col.y + micro_helio_id / row_col.x] = true;
        atomicAdd(&(d_sb_result[helio_id].x), origin.x);
        atomicAdd(&(d_sb_result[helio_id].y), origin.y);
        atomicAdd(&(d_sb_result[helio_id].z), origin.z);
        atomicAdd(&(d_sb_result[helio_id].w), 1);
    }
    //    rectGridDDA::test(sun_dir);
}

__global__ static void shadowBlockingWithSharedMemory(RectangleGrid rectangle_grid, float3* d_verts,
                                                      float4* d_sb_result, bool* d_microhelio,
                                                      int2 row_col, int n, float3 sun_dir,
                                                      int num_threads)
{
    long long my_id = getThreadID();
    if (my_id >= row_col.x * row_col.y * n)
    {
        return;
    }
    // size = row_col.x
    extern __shared__ float4 temp[];
    const int row_col_mul = row_col.x * row_col.y;
    int helio_id = my_id / (row_col_mul);
    int micro_helio_id = my_id % (row_col_mul);
    //    printf("%d\n",micro_helio_id);
    const auto& vert0 = d_verts[helio_id * 3];
    const auto& vert1 = d_verts[helio_id * 3 + 1];
    const auto& vert2 = d_verts[helio_id * 3 + 2];
    const auto vec1 = vert1 - vert0;
    const auto vec2 = vert2 - vert1;
    float3 normal = normalize(cross(vec2, vec1));
    // if (dot(normal, sun_dir) > 0)
    // {
    //     normal = -normal;
    // }
    normal -= 2 * normal * float(dot(normal, sun_dir) > 0);
    float3 origin = (0.5F + micro_helio_id / row_col.y) * (vec2) / row_col.x +
                    (0.5F + micro_helio_id % row_col.y) * (vec1) / row_col.y + vert0;
    //    printf("%f %f %f\n",origin.x,origin.y,origin.z);
    float3 reflect_dir = normalize(reflect(sun_dir, normal));
    sun_dir = -sun_dir;
    bool flag_collision = collision(origin, sun_dir, &rectangle_grid, d_verts, helio_id) ||
                          collision(origin, reflect_dir, &rectangle_grid, d_verts, helio_id);
    d_microhelio[micro_helio_id + helio_id * row_col.x * row_col.y] = flag_collision;
    // atomicAdd(&(d_sb_result[helio_id].x), origin.x);
    // atomicAdd(&(d_sb_result[helio_id].y), origin.y);
    // atomicAdd(&(d_sb_result[helio_id].z), origin.z);
    // atomicAdd(&(d_sb_result[helio_id].w), 1);
    temp[threadIdx.x].x = bit::conditionalAdd2(flag_collision, 0.0F, origin.x);
    temp[threadIdx.x].y = bit::conditionalAdd2(flag_collision, 0.0F, origin.y);
    temp[threadIdx.x].z = bit::conditionalAdd2(flag_collision, 0.0F, origin.z);
    temp[threadIdx.x].w = static_cast<float>(flag_collision);
    __syncthreads();
    if (threadIdx.x == 0)
    {
        float4 partial_sum = {0, 0, 0, 0};
        // const auto total_threads_per_grid = std::gcd(row_col.x, row_col.y);
        for (int i = 0; i < num_threads; i++)
        {
            partial_sum += temp[i];
        }
        atomicAdd(&(d_sb_result[helio_id].x), partial_sum.x);
        atomicAdd(&(d_sb_result[helio_id].y), partial_sum.y);
        atomicAdd(&(d_sb_result[helio_id].z), partial_sum.z);
        atomicAdd(&(d_sb_result[helio_id].w), partial_sum.w);
    }
    //    rectGridDDA::test(sun_dir);
}

__global__ static void shadowBlockingWithSharedMemory2(RectangleGrid rectangle_grid,
                                                       float3* d_verts, float4* d_sb_result,
                                                       bool* d_microhelio, int2 row_col, int n,
                                                       float3 sun_dir)
{
    const auto my_id = getThreadID();
    const int row_col_mul = row_col.x * row_col.y;
    if (my_id >= row_col_mul * n)
    {
        return;
    }
    // size = row_col.x
    extern __shared__ float4 temp[];
    int helio_id = my_id / (row_col_mul);
    int micro_helio_id = my_id % (row_col_mul);
    //    printf("%d\n",micro_helio_id);
    const auto vert0 = d_verts[helio_id * 3];
    const auto vert1 = d_verts[helio_id * 3 + 1];
    const auto vert2 = d_verts[helio_id * 3 + 2];
    const auto vec1 = vert1 - vert0;
    const auto vec2 = vert2 - vert1;
    float3 normal = normalize(cross(vec2, vec1));
    // if (dot(normal, sun_dir) > 0)
    // {
    //     normal = -normal;
    // }
    normal -= 2 * normal * float(dot(normal, sun_dir) > 0);
    float3 origin = (0.5F + micro_helio_id / row_col.y) * (vec2) / row_col.x +
                    (0.5F + micro_helio_id % row_col.y) * (vec1) / row_col.y + vert0;
    //    printf("%f %f %f\n",origin.x,origin.y,origin.z);
    float3 reflect_dir = normalize(reflect(sun_dir, normal));
    sun_dir = -sun_dir;
    bool flag_collision = collision(origin, sun_dir, &rectangle_grid, d_verts, helio_id) ||
                          collision(origin, reflect_dir, &rectangle_grid, d_verts, helio_id);
    d_microhelio[micro_helio_id + helio_id * row_col_mul] = flag_collision;
    atomicAdd(&(temp[0].x), bit::conditionalAdd2(flag_collision, 0.0F, origin.x));
    atomicAdd(&(temp[0].y), bit::conditionalAdd2(flag_collision, 0.0F, origin.y));
    atomicAdd(&(temp[0].z), bit::conditionalAdd2(flag_collision, 0.0F, origin.z));
    atomicAdd(&(temp[0].w), bit::conditionalAdd2(flag_collision, 0.0F, 1.0F));
    __syncthreads();
    if (threadIdx.x == 0)
    {
        atomicAdd(&(d_sb_result[helio_id].x), temp[0].x);
        atomicAdd(&(d_sb_result[helio_id].y), temp[0].y);
        atomicAdd(&(d_sb_result[helio_id].z), temp[0].z);
        atomicAdd(&(d_sb_result[helio_id].w), temp[0].w);
    }
    //    rectGridDDA::test(sun_dir);
}

__global__ static void shadowBlocking2(std::tuple<int3, int*, int*, size_t> rectangle_grid_params,
                                       float3* d_verts, float4* d_sb_result, bool* d_microhelio,
                                       int2 row_col, int n, float3 sun_dir)
{
    long long my_id = getThreadID();
    if (my_id >= row_col.x * row_col.y * n)
    {
        return;
    }
    RectangleGrid rectangle_grid;
    rectangle_grid.subgrid_num_ = std::get<0>(rectangle_grid_params);
    rectangle_grid.d_grid_helio_match_ = std::get<1>(rectangle_grid_params);
    rectangle_grid.d_grid_helio_index_ = std::get<2>(rectangle_grid_params);
    rectangle_grid.num_grid_helio_match_ = std::get<3>(rectangle_grid_params);
    int helio_id = my_id / (row_col.x * row_col.y);
    int micro_helio_id = my_id % (row_col.x * row_col.y);
    //    printf("%d\n",micro_helio_id);
    auto vert0 = d_verts[helio_id * 3];
    auto vert1 = d_verts[helio_id * 3 + 1];
    auto vert2 = d_verts[helio_id * 3 + 2];
    float3 normal = normalize(cross(vert2 - vert1, vert1 - vert0));
    if (dot(normal, sun_dir) > 0)
    {
        normal = -normal;
    }
    float3 origin = (0.5 + micro_helio_id / row_col.y) * (vert2 - vert1) / row_col.x +
                    (0.5 + micro_helio_id % row_col.y) * (vert1 - vert0) / row_col.y + vert0;
    //    printf("%f %f %f\n",origin.x,origin.y,origin.z);
    float3 reflect_dir = normalize(reflect(sun_dir, normal));
    sun_dir = -sun_dir;
    if (collision(origin, sun_dir, &rectangle_grid, d_verts, helio_id) ||
        collision(origin, reflect_dir, &rectangle_grid, d_verts, helio_id))
    {
        d_microhelio[micro_helio_id + helio_id * row_col.x * row_col.y] = true;
        atomicAdd(&(d_sb_result[helio_id].x), origin.x);
        atomicAdd(&(d_sb_result[helio_id].y), origin.y);
        atomicAdd(&(d_sb_result[helio_id].z), origin.z);
        atomicAdd(&(d_sb_result[helio_id].w), 1);
    }
    //    rectGridDDA::test(sun_dir);
}

void ShadowBlockCalculator::calculate()
{
    if (this->rectangle_grid_ == nullptr)
    {
        const auto error_msg =
            solar::format("In function {}: rectangle_grid_ is nullptr", PRETTY_FUNCTION_NAME);
        solar::println(error_msg);
        throw error_msg;
    }

    copyHeliostatVerticesToGPU();
    initResultArray();
    // 默认整个镜场在同一个grid中
    // 假设所有定日镜的尺寸都相同
    // auto helio_size = scene->getHeliostats()[0]->getSize();
    //    int m = int(helio_size.x/helios_pixel_size+0.5);
    //    int n = int(helio_size.z/helios_pixel_size+0.5);
    //    const int m = 10, n = 10;
    const int m = helio_row_col_.x;
    const int n = helio_row_col_.y;
    // const auto micro_helio_num_foreach_helio = m * n;
    // 离散为10×10
    //    m = 10; n = 10;

    // const int n_threads = 512;
    int n_threads = helio_row_col_.x;
    if (m * n < 512)
    {
        n_threads = m * n;
    }
    // const int n_threads = std::gcd(helio_row_col_.x, helio_row_col_.y);
    dim3 n_blocks;
    setThreadBlocks(n_blocks, n_threads, m * n * n_helio_, true);
    //    printf("%lld\n", heliostatArgument.numberOfMicroHeliostats *
    //    sunrayArgument.numberOfLightsPerGroup);

    // Runtime failure, a Clang bug ? "kernel启动失败"
    // shadowBlocking<<<n_blocks, n_threads>>>(*rectangle_grid_, d_verts_, d_sb_result,
    // d_microhelio_,
    //                                         make_int2(m, n), n_helio_, sun_direction);

    // cudaStreamSynchronize(rectangle_grid_->stream_);
    rectangle_grid_->syncCudaStream();

    CHECK_CUDA_ERRORS(cudaStreamSynchronize(this->stream_));
    // shadowBlocking2<<<n_blocks, n_threads, 0, stream_>>>(
    //     {rectangle_grid_->subgrid_num_, rectangle_grid_->d_grid_helio_match_,
    //      rectangle_grid_->d_grid_helio_index_, rectangle_grid_->num_grid_helio_match_},
    //     d_verts_, d_sb_result_, d_microhelio_, make_int2(m, n), n_helio_, sun_direction_);
    shadowBlocking<<<n_blocks, n_threads, 0, stream_>>>(*rectangle_grid_, d_verts_, d_sb_result_,
                                                        d_microhelio_, d_microhelio_per_row_,
                                                        make_int2(m, n), n_helio_, sun_direction_);
    // shadowBlockingWithSharedMemory2<<<n_blocks, n_threads, sizeof(float4) * 1, stream_>>>(
    //     *rectangle_grid_, d_verts_, d_sb_result_, d_microhelio_, make_int2(m, n), n_helio_,
    //     sun_direction_);

    // CHECK_CUDA_ERRORS(cudaDeviceSynchronize());
    CHECK_CUDA_ERRORS(cudaGetLastError());
}

void ShadowBlockCalculator::calculateCenterBias(float3* d_heliostat_positions)
{
    thrust::device_ptr<float3> heliostat_positions_ptr(d_heliostat_positions);
    thrust::device_ptr<float4> result_ptr(this->d_sb_result_);
    const auto helio_row_col_mul = this->helio_row_col_.x * this->helio_row_col_.y;
    auto zip_iterator = thrust::make_zip_iterator(heliostat_positions_ptr, result_ptr);
    auto result = thrust::async::transform(
        thrust::device.on(stream_), zip_iterator, zip_iterator + this->n_helio_, result_ptr,

        [helio_row_col_mul] __host__ __device__(const thrust::tuple<float3, float4>& iter)
        {
            auto position = thrust::get<0>(iter);
            auto current_result = thrust::get<1>(iter);
            float3 center_bias =
                (position * helio_row_col_mul - structCopy<float3>(current_result)) /
                (helio_row_col_mul - current_result.w);
            auto shadow_block_factor = current_result.w / helio_row_col_mul;
            return make_float4(center_bias, shadow_block_factor);
        });
}
